# 2023/7/23 13:00
# xiashuobad
# transform_types
# PyCharm
"""
 * _ooOoo_
 * o8888888o
 * 88" . "88
 * (| -_- |)
 *  O\ = /O
 * ___/`---'\____
 * .   ' \\| |// `.
 * / \\||| : |||// \
 * / _||||| -:- |||||- \
 * | | \\\ - /// | |
 * | \_| ''\---/'' | |
 * \ .-\__ `-` ___/-. /
 * ___`. .' /--.--\ `. . __
 * ."" '< `.___\_<|>_/___.' >'"".
 * | | : `- \`.;`\ _ /`;.`/ - ` : | |
 * \ \ `-. \_ __\ /__ _/ .-` / /
 * ======`-.____`-.___\_____/___.-`____.-'======
 * `=---=' bug泛滥 佛已瘫痪
"""
import numpy as np
# import torch
from affine import Affine
# from mmcv import is_list_of
from osgeo import gdalconst, osr
from typing import List, Union
from collections import OrderedDict
from pathlib import Path
from constant import DICT_GDAL_DATA_TYPE, NEW_NORM_CFG
from misc import is_list_of


def pixel_geo_transform(x_pixel, y_pixel, geo_transform) -> tuple:
    """
    像素坐标和地理坐标仿射变换
    :param x_pixel: 像素坐标x
    :param y_pixel: 像素坐标y
    :param geo_transform:像素坐标(0,0)对应的坐标变换系数: (GT(0),GT(1),GT(2),GT(3),GT(4),GT(5))
        GT(0) 左上像素左上角的x坐标。
        GT(1) w-e像素分辨率/像素宽度。
        GT(2) 行旋转（通常为零）。
        GT(3) 左上像素左上角的y坐标。
        GT(4) 列旋转（通常为零）。
        GT(5) n-s像素分辨率/像素高度（北上图像为负值）。
    :return: 像素坐标(x_pixel,y_pixel) 对应的坐标变换系数
    """
    new_geo_transform = np.array(geo_transform)
    x_geo, y_geo = pixel2geo(x_pixel, y_pixel, geo_transform)
    new_geo_transform[[0, 3]] = x_geo, y_geo
    return tuple(new_geo_transform)


def adjust_transform(geo_transform) -> tuple:
    """
    当原始影像没有坐标系，而为默认(0,1,0,0,0,1)时，自动调整，上下翻转
    Args:
        geo_transform: 坐标

    Returns:

    """
    geo_transform = geo_transform2gdal(geo_transform)  # 先转换成gdal格式
    if geo_transform == (0, 1, 0, 0, 0, 1):  # 这是默认投影，会导致上下翻转
        # proj.ImportFromEPSG(4326)
        geo_transform = (0, 1, 0, 0, 0, -1)

    return geo_transform


def pixel2geo(x_pixel, y_pixel, geo_transform):
    """
    像素坐标转地理坐标
    Args:
        point:像素坐标(x,y)
        geo_transform:坐标变换系数 (x_left,dx,r_rotation,y_top,c_rotation,dy)

    Returns:

    """
    geo_transform = geo_transform2gdal(geo_transform)
    x_geo = geo_transform[1] * x_pixel + y_pixel * geo_transform[2] + geo_transform[0]
    y_geo = geo_transform[4] * x_pixel + y_pixel * geo_transform[5] + geo_transform[3]
    return x_geo, y_geo


def geo2pixel(x_geo, y_geo, geo_transform):
    """
    地理坐标转像素坐标
    Args:
        point:(x_geo,y_geo)
        geo_transform:

    Returns:

    """
    # geo_transform = geo_transform2gdal(geo_transform)
    # x_pixel = (geo_transform[5] * x_geo - geo_transform[2] * y_geo + geo_transform[2] * geo_transform[3] -
    #            geo_transform[0] * geo_transform[5]) / (
    #                   geo_transform[1] * geo_transform[5] - geo_transform[2] * geo_transform[4])
    # y_pixel = (geo_transform[1] * y_geo - geo_transform[4] * x_geo + geo_transform[0] * geo_transform[4] -
    #            geo_transform[1] * geo_transform[3]) / (
    #                   geo_transform[1] * geo_transform[5] - geo_transform[2] * geo_transform[4])

    a = np.array([[geo_transform[1], geo_transform[2]], [geo_transform[4], geo_transform[5]]])
    b = np.array([x_geo - geo_transform[0], y_geo - geo_transform[3]])
    x_pixel, y_pixel = np.linalg.solve(a, b)  # 使用numpy的linalg.solve进行二元一次方程的求解
    return round(abs(x_pixel)), round(abs(y_pixel))


def geo_coord2proj_coord(point: (tuple, list), dataset):
    """
    地理坐标转投影坐标
    Args:
        point:(x_geo,y_geo)
        dataset:

    Returns:

    """
    prosrs, geosrs = get_spatial_reference(dataset)
    ct = osr.CoordinateTransformation(geosrs, prosrs)
    x_geo, y_geo = point
    x_proj, y_proj, z_proj = ct.TransformPoint(x_geo, y_geo)
    return x_proj, y_proj


def proj_coord2geo_coord(point: (tuple, list), dataset):
    """
    投影坐标转地理坐标
    Args:
        point:(x_proj,y_proj)
        dataset:

    Returns:

    """
    prosrs, geosrs = get_spatial_reference(dataset)
    ct = osr.CoordinateTransformation(prosrs, geosrs)
    x_proj, y_proj = point
    x_geo, y_geo, z_geo = ct.TransformPoint(x_proj, y_proj)
    return x_geo, y_geo


def get_spatial_reference(dataset):
    """
    获取影像的投影和地理坐标系
    Args:
        dataset:

    Returns:

    """
    prosrs = osr.SpatialReference()
    prosrs.ImportFromWkt(dataset.GetProjection())
    geosrs = prosrs.CloneGeogCS()
    return prosrs, geosrs


def geo_transform2affine(geo_transform: tuple):
    """
    转换transform为affine格式
    :param geo_transform:
    :return:
    """
    if not isinstance(geo_transform, Affine):  # 当地理坐标仿射变换是gdal格式时，转换为affine格式
        assert isinstance(geo_transform, tuple) and len(geo_transform) == 6, '只支持gdal格式(x,dx,rx,y,ry,dy)的转换affine格式'
        geo_transform = Affine.from_gdal(*geo_transform)
    return geo_transform


def geo_transform2gdal(geo_transform: Affine):
    """
    转换transform为gdal格式
    :param geo_transform:
    :return:
    """
    if isinstance(geo_transform, Affine):
        geo_transform = Affine.to_gdal(geo_transform)
    return geo_transform


def reshape_as_image(img_array: np.ndarray):
    """
    将图片array转换为正常图片格式，rgb通道在第三维
    Args:
        img_array:
    Returns:
    """

    assert (length_shape := len(img_array.shape)) in {2, 3}, 'img_array输入格式错误！'
    if length_shape == 3 and np.argmin(
            img_array.shape) == 0:  # 这种写法只能判断尺寸大于3*3的影像array的rgb通道在哪个维度，如果在第1个维度，则进行转换
        img_array = img_array.transpose((1, 2, 0))
    return img_array


def reshape_as_raster(img_array: np.ndarray):
    """
    将图片array转换为栅格格式，rgb通道在第1维
    Args:
        img_array:
    Returns:
    """

    assert (length_shape := len(img_array.shape)) in {2, 3}, 'img_array输入格式错误！'
    if length_shape == 3 and np.argmin(
            img_array.shape) == 2:  # 这种写法只能判断尺寸大于3*3的影像array的rgb通道在哪个维度，如果在第3个维度，则进行转换
        img_array = img_array.transpose((2, 0, 1))
    return img_array


def color_code2rgb(color_code: str) -> tuple:
    """
    颜色代码转rgb元组
    :param color_code:
    :return:
    """
    if isinstance(color_code, (list, tuple)) and len(color_code) == 3:
        return color_code
    assert isinstance(color_code, str), 'color_code参数类型错误！'
    color_code = color_code.strip('#')
    assert len(color_code) == 6, '颜色代码位数应为16进制6位数！'
    rgb = (int(color_code[:2], 16), int(color_code[2:4], 16), int(color_code[4:], 16))
    return rgb


def format_colormap2rasterio(color_mapping: List[Union[int, list, tuple, str]]) -> dict:
    """
    将PIL格式的rgb映射格式转换成rasterio的格式
    Args:
        color_mapping:

    Returns:

    """
    if isinstance(color_mapping, dict):
        if is_list_of(list(color_mapping.values()), tuple) or is_list_of(list(color_mapping.values()), list):
            return color_mapping
        else:
            for key, value in color_mapping.items():
                color_mapping[key] = color_code2rgb(value)
            return color_mapping
    assert isinstance(color_mapping, (list, tuple)), 'color_mapping参数类型错误！'
    if isinstance(color_mapping[0], (int, list, tuple)):
        color_mapping = np.array(color_mapping).reshape((-1, 3))
        color_mapping = {i: tuple(rgb) for i, rgb in enumerate(color_mapping)}
    elif isinstance(color_mapping[0], str):
        color_mapping = {color_code2rgb(color_code) for i, color_code in enumerate(color_mapping)}
    else:
        raise Exception('color_mapping参数类型错误！')
    return color_mapping


def format_colormap2pil(color_mapping: (list, dict)) -> List[int]:
    """
    将rasterio格式的rgb映射表转换成PIL格式
    :param color_mapping:
    :return:
    """
    if isinstance(color_mapping, list):
        if is_list_of(color_mapping, int):
            return color_mapping
        assert is_list_of(color_mapping, tuple) or is_list_of(color_mapping, list)
        color_mapping = list(np.array(color_mapping).flatten())
        return color_mapping
    assert isinstance(color_mapping, dict) and is_list_of(list(color_mapping.keys()), int) and is_list_of(
        list(color_mapping.values()), tuple), 'color_mapping应为rasterio风格的字典格式！'
    res_color_mapping = [0, 0, 0] * (max(color_mapping.keys()) + 1)
    for k, v in color_mapping.items():
        res_color_mapping[3 * k:3 * k + 3] = v
    return res_color_mapping


def dtype2gdal_type(dtype: str):
    """
    dtype转gdal数据类型
    """
    return DICT_GDAL_DATA_TYPE.get(dtype, gdalconst.GDT_Byte)


# def model_out2pred(logit: (torch.Tensor, np.ndarray)) -> np.ndarray:
#     """
#     从网络最后一层输出得概率图转换到预测图
#     :param logit: 网络最后一层输出结果
#     :return: 返回预测图
#     """
#     if isinstance(logit, torch.Tensor):
#         logit = logit.cpu().detach().numpy()
#     assert len(logit.shape) == 4, 'logit的shape格式必须为模型输出格式：(b,c,h,w)!'
#     out_channels = logit.shape[1]
#     if out_channels == 1:  # 二分类只输出1个通道的情况
#         y_pred = 1 / (1 + np.exp(-logit))
#         y_pred = np.where(y_pred > 0.5, 1, 0).astype(np.uint8)
#         y_pred = y_pred.squeeze(1)
#     else:  # 多分类（包括二分类）
#         y_pred = logit.argmax(axis=1).astype(np.uint8)
#     return y_pred


# def img_array2input_tensor(image_array: np.ndarray) -> torch.Tensor:
#     """
#      模型输入图片array转tensor，
#      用于网络输入的格式:
#         image: (batch_size,channel,h,w)
#     :param image_array:
#     :return:
#     """
#     if len(image_array.shape) == 2:
#         # image_array = image_array[None, :, :]
#         image_array = np.stack((image_array,) * 3, axis=0)
#     image_array = reshape_as_raster(image_array)
#     image_tensor = torch.as_tensor(image_array, dtype=torch.float32)
#     image_tensor = image_tensor.unsqueeze(0)
#     return image_tensor


def format_class_info(list_class_info: (list, tuple)):
    """
    返回格式化后的class_names: 类似[{'class_id':0,'一级类别':'道路','二级类别':'铁路','jp_bm':666666},{'class_id':1,...},...]
    Args:
        list_class_info:

    Returns:

    """
    assert isinstance(list_class_info, (list, tuple))
    class_names = list(list_class_info)
    class_fields = {}  # 初始化类别字段列表
    for class_info in class_names:  # 遍历获取所有的类别字段
        if isinstance(class_info, dict):
            for field, value in class_info.items():
                class_fields[field] = type(value)
    if len(class_fields) == 0:  # 未设置类别字段时，默认添加 'class_name'字段
        class_fields['class_name'] = str
    for i, class_info in enumerate(class_names):
        new_class_info = {'class_id': i}
        if isinstance(class_info, str):  # 默认使用第一个类别字段，其他字段的value设置为'###'
            for j, field in enumerate(class_fields):
                if j == 0:
                    new_class_info[field] = class_info
                else:
                    new_class_info[field] = new_class_info.get(field, class_fields[field](0))
        elif isinstance(class_info, dict):
            for field, field_type in class_fields.items():
                new_class_info[field] = class_info.get(field, field_type(0))
        else:
            raise ValueError(f'{class_info}: 数据格式有误！')

        class_names[i] = new_class_info
    class_fields['class_id'] = int
    return {'class_names': class_names, 'class_fields': class_fields}


def format_models_info(models_info: (str, list, tuple), num_class: int, norm_cfg: dict = None) -> list:
    """
    格式化models_info
    Args:
        models_info:
        num_class: 总类别数，包括背景
        norm_cfg: 归一化参数

    Returns:

    """
    if isinstance(models_info, (str, Path, tuple, dict)):
        models_info = [models_info]
    models_info = list(models_info)
    for i, model_info in enumerate(models_info):
        models_info[i] = OrderedDict(
            {'model_path': '',
             'class_ids': list(range(1, num_class)),
             'norm_cfg': norm_cfg or NEW_NORM_CFG,
             'model_name': '',
             'min_filter_area': 100,
             'min_filter_size': 5000,
             'cfg_path': None
             }
        )
        if isinstance(model_info, (str, Path)):
            models_info[i]['model_path'] = model_info
        elif isinstance(model_info, (list, tuple)):
            if len(model_info) >= 1:
                models_info[i]['model_path'] = model_info[0]
            if len(model_info) >= 2:
                if isinstance(class_ids := model_info[1], int):
                    class_ids = [class_ids]
                models_info[i]['class_ids'] = class_ids
            if len(model_info) >= 3:
                models_info[i]['norm_cfg'] = model_info[2]
            if len(model_info) >= 4:
                models_info[i]['model_name'] = model_info[3]
            if len(model_info) >= 5:
                models_info[i]['min_filter_area'] = model_info[4]
            if len(model_info) >= 6:
                models_info[i]['cfg_path'] = model_info[5]
        elif isinstance(model_info, dict):
            models_info[i].update(model_info)
        else:
            raise ValueError(f'{model_info}: 数据格式有误！')

    return models_info


def format_model_dict(model_dict: dict) -> dict:
    """
    格式化模型字典
    Args:
        model_dict: 模型字典，同测试配置文件

    Returns:

    """
    new_model_dict = OrderedDict()
    class_info = format_class_info(model_dict['class_info'])
    models_info = format_models_info(model_dict['models_info'], len(class_info['class_names']),
                                     model_dict.get('norm_cfg', None))
    new_model_dict['models_info'] = models_info
    new_model_dict['model_merge_mode'] = model_dict.get('model_merge_mode', 'last')
    new_model_dict['class_info'] = class_info
    new_model_dict['color_map'] = model_dict.get('color_map', None)
    new_model_dict['shp_type'] = model_dict.get('shp_type', '面')
    new_model_dict['tta'] = model_dict.get('tta', None)
    new_model_dict['input_size'] = model_dict.get('input_size', (512, 512))
    new_model_dict['threshold'] = model_dict.get('threshold', 0.3)
    new_model_dict['simplify_tolerance'] = model_dict.get('simplify_tolerance', 2)

    return new_model_dict
